import requests
import json
import re
import base64
from datetime import datetime, timedelta

API_KEY = "sk-*"  
API_URL = "https://api.deepseek.com/v1/chat/completions"

HEADERS = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {API_KEY}"
}

# helper function 1: encode image to base64
def encode_image_to_base64(filepath):
    try:
        with open(filepath, "rb") as f:
            encoded = base64.b64encode(f.read()).decode("utf-8")
        return encoded
    except FileNotFoundError:
        print(f"[错误] 图片路径不存在: {filepath}")
        return ""

# helper function 2: extract json from response
def extract_json_from_response(response_text):
    match = re.search(r"```json\s*(.*?)\s*```", response_text, re.DOTALL)
    if match:
        json_str = match.group(1)
    else:
        json_str = response_text
    try:
        return json.loads(json_str)
    except json.JSONDecodeError as e:
        print("无法解析 JSON：", e)
        print("原始内容：\n", response_text)
        return None
    

##### step 1: reformat and filter data into json
#####

def filter_and_format_messages(raw_messages):
    filtered_msgs = []
    for msg in raw_messages:
        if msg["type"] == "text":
            text = msg["content"]
            if any(x in text for x in ["收到", "谢谢", "了解了"]):
                continue
            filtered_msgs.append({
                "type": "text",
                "content": text,
                "timestamp": msg.get("timestamp", "")
            })
        elif msg["type"] == "image":
            filtered_msgs.append({
                "type": "image",
                "content": {
                    "base64": msg["content"]["base64"],
                    "filename": msg["content"]["filename"]
                },
                "timestamp": msg.get("timestamp", "")
            })
    return filtered_msgs


##### step 2: cluster filtered messages into events
#####

def cluster_filtered_robot_events(messages):
    """仅使用文字进行聚类，不传图片"""
    CLUSTER_PROMPT = """
你是一个结构化助手。请将以下巡检文字信息聚类成多个事件卡。每个事件卡包括：

1. 一个简要的聚合描述；
2. 所有属于该事件的原始留言（保留内容）；

请返回 JSON 格式，形如：
[
  {{
    "event_id": 1,
    "summary": "在3号楼东南角发现脚手架松动",
    "messages": [
      {{"type": "text", "content": "...", "timestamp": "..."}},
      ...
    ]
  }},
  ...
]

以下是巡检留言信息：
{records}
"""
    text_messages = [msg for msg in messages if msg["type"] == "text"]

    lines = []
    for msg in text_messages:
        ts = f"[{msg['timestamp']}]" if msg.get("timestamp") else ""
        lines.append(f"- {ts} [TEXT] {msg['content']}")

    prompt = CLUSTER_PROMPT.format(records="\n".join(lines))

    payload = {
        "model": "deepseek-chat",
        "messages": [
            {"role": "system", "content": "你是一个巡检机器人助手"},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.3
    }

    response = requests.post(API_URL, headers=HEADERS, json=payload)
    if response.status_code == 200:
        return extract_json_from_response(response.json()["choices"][0]["message"]["content"])
    else:
        print("请求失败:", response.status_code, response.text)
        return None


def assign_images_to_events(events, images, time_window_minutes=15):
    for event in events:
        event_time_range = [
            datetime.fromisoformat(msg["timestamp"].replace("Z", "+00:00"))
            for msg in event["messages"]
        ]
        min_time = min(event_time_range)
        max_time = max(event_time_range)
        window_start = min_time - timedelta(minutes=time_window_minutes)
        window_end = max_time + timedelta(minutes=time_window_minutes)

        # 找所有在这个时间窗内的图片
        candidate_imgs = []
        for img in images:
            try:
                img_time = datetime.fromisoformat(img["timestamp"].replace("Z", "+00:00"))
                if window_start <= img_time <= window_end:
                    candidate_imgs.append({
                        "base64": img["content"]["base64"],
                        "filename": img["content"]["filename"],
                        "timestamp": img["timestamp"]
                    })
            except:
                continue

        event["candidate_images"] = candidate_imgs
    return events


def main():
    print("Running clustering pipeline...\n")

    img2_base64 = encode_image_to_base64("/Users/wangzijian/Desktop/实习gauz/gauz_ai_clustering/IMG_2.jpg")
    img3_base64 = encode_image_to_base64("/Users/wangzijian/Desktop/实习gauz/gauz_ai_clustering/image03.jpg")
    # 示例消息结构
    raw_messages = [
        {
            "type": "text",
            "content": "在3号楼东南角发现脚手架松动，有坠落风险",
            "timestamp": "2024-04-01T08:30:00Z"
        },
        {
            "type": "text",
            "content": "责任人张三已通知整改",
            "timestamp": "2024-04-01T08:31:10Z"
        },
        {
            "type": "text",
            "content": "收到，谢谢",
            "timestamp": "2024-04-01T08:35:00Z"
        },
        {
            "type": "text",
            "content": "2号楼发现电缆裸露，存在触电危险",
            "timestamp": "2024-04-01T10:00:00Z"
        },
        {
            "type": "image",
            "content": {
                "base64": img2_base64,
                "filename": "img_cable_02.jpg"
            },
            "timestamp": "2024-04-01T10:01:00Z"
        },
        {
            "type": "text",
            "content": "监理回复：了解了",
            "timestamp": "2024-04-01T10:05:00Z"
        },
        {
            "type": "image",
            "content": {
                "base64": img3_base64,
                "filename": "img_cable_03.jpg"
            },
            "timestamp": "2024-04-01T11:05:00Z"
        }
    ]

    # step 1
    filtered_msgs = filter_and_format_messages(raw_messages)
    text_msgs = [m for m in filtered_msgs if m["type"] == "text"]
    image_msgs = [m for m in filtered_msgs if m["type"] == "image"]

    # step 2
    clustered_events = cluster_filtered_robot_events(text_msgs)

    # step 3
    clustered_with_images = assign_images_to_events(clustered_events, image_msgs, time_window_minutes=10)

    # step 4
    print("最终事件卡：")
    for event in clustered_with_images:
        print(f"\nEvent {event['event_id']} Summary: {event['summary']}")
        print(f"  Messages: {(event['messages'])}")
        print(f"  Candidate Images: {len(event['candidate_images'])}")

if __name__ == "__main__":
    main()
